"""
数据库工具模块
提供数据库连接和SQL查询功能
"""
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
import dmPython
import re

# 加载环境变量
load_dotenv()

def get_database_connection():
    """
    获取达梦数据库连接
    """
    try:
        # 创建数据库连接
        conn = dmPython.connect(
            user=os.getenv('DM_USER', 'DPA'), 
            password=os.getenv('DM_PASSWORD', 'Szzt#2023'), 
            server=os.getenv('DM_HOST', '10.255.102.23'),
            port=int(os.getenv('DM_PORT', 5236)), 
            autoCommit=True
        )
        return conn
    except Exception as e:
        print(f"达梦数据库连接失败: {e}")
        return None

class DMDatabase:
    """达梦数据库适配器类，用于与LangChain集成"""
    
    def __init__(self, connection=None):
        if connection is None:
            connection = get_database_connection()
        self.connection = connection
    
    @property
    def dialect(self):
        """返回数据库方言，模拟SQLDatabase的dialect属性"""
        return "dm"  # 达梦数据库方言
    
    def get_usable_table_names(self):
        """获取可用的表名"""
        try:
            cursor = self.connection.cursor()
            # 达梦数据库查询表的SQL语句
            cursor.execute("SELECT TABLE_NAME FROM USER_TABLES")
            tables = [row[0] for row in cursor.fetchall()]
            cursor.close()
            return tables
        except Exception as e:
            print(f"获取表名失败: {e}")
            return []
    
    def run(self, command, fetch="all"):
        """执行SQL命令"""
        try:
            # 清理SQL命令中的反引号（达梦数据库不支持）
            command = command.replace('`', '')
            # 移除可能存在的多余空白字符
            command = command.strip()
            
            # 移除可能存在的注释符号（达梦数据库可能不支持某些注释格式）
            command = re.sub(r'--.*', '', command)
            
            # 特殊处理：移除可能导致解析错误的别名表达式（加强版）
            command = re.sub(r'\[([A-Z_]+\.[A-Z_]+)\]', r'\1', command)
            
            cursor = self.connection.cursor()
            cursor.execute(command)
            
            if fetch == "all":
                result = cursor.fetchall()
            elif fetch == "one":
                result = cursor.fetchone()
            else:
                result = None
            
            cursor.close()
            return str(result)
        except Exception as e:
            raise Exception(f"执行SQL失败: {e}")
    #获取表构建信息（表名，字段名，数据类型）输出格式
    def get_table_info(self, table_names=None):
        """获取表信息CREATE TABLE table_name (
    column1 datatype -- comment,
    column2 datatype -- comment
  )"""
        if table_names is None:
            table_names = self.get_usable_table_names()
        
        template = """
            CREATE TABLE {table_name} (
            {columns}
            )
        """
        
        table_info = []
        try:
            cursor = self.connection.cursor()
            for table_name in table_names:
                # 查询表结构信息，包括列注释
                # 使用USER_TAB_COLUMNS视图来获取列信息，确保列名正确
                cursor.execute(f"""
                    SELECT COLUMN_NAME, DATA_TYPE, (SELECT COMMENTS FROM USER_COL_COMMENTS WHERE TABLE_NAME = '{table_name}' AND COLUMN_NAME = UTC.COLUMN_NAME) AS COMMENTS
                    FROM USER_TAB_COLUMNS UTC
                    WHERE TABLE_NAME='{table_name}'
                    ORDER BY COLUMN_ID
                """)
                columns = cursor.fetchall()
                
                column_info = []
                for column in columns:
                    column_name, data_type, comments = column[0], column[1], column[2]
                    column_desc = f"  {column_name} {data_type}"
                    if comments:
                        column_desc += f" -- {comments}"  # 使用SQL标准注释语法
                    column_info.append(column_desc)
                
                table_info.append(template.format(
                    table_name=table_name,
                    columns=",\n".join(column_info)
                ))
            cursor.close()
            
            return "\n\n".join(table_info)
        except Exception as e:
            return f"获取表信息失败: {e}"
    
    def get_table_comments(self, table_names=None):
        """获取表的注释信息"""
        if table_names is None:
            table_names = self.get_usable_table_names()
        
        try:
            cursor = self.connection.cursor()
            table_comments = {}
            
            for table_name in table_names:
                # 查询表注释
                cursor.execute(f"""
                    SELECT COMMENTS 
                    FROM USER_TAB_COMMENTS 
                    WHERE TABLE_NAME='{table_name}'
                """)
                result = cursor.fetchone()
                if result and result[0]:
                    table_comments[table_name] = result[0]
                    
            cursor.close()
            return table_comments
        except Exception as e:
            print(f"获取表注释失败: {e}")
            return {}

    def get_all_table_and_column_info(self):
        """
        获取所有表和列的信息，包括注释
        返回格式化的字符串，用于自然语言处理
        """
        try:
            # 获取所有表名
            table_names = self.get_usable_table_names()
            
            # 获取所有表的注释
            table_comments = self.get_table_comments(table_names)
            
            # 构建信息字符串
            info_parts = []
            
            for table_name in table_names:
                info_part = f"表名: {table_name}\n"
                if table_name in table_comments:
                    info_part += f"表注释: {table_comments[table_name]}\n"
                
                # 获取表的列信息
                cursor = self.connection.cursor()
                # 使用USER_TAB_COLUMNS视图来获取列信息
                cursor.execute(f"""
                    SELECT COLUMN_NAME, DATA_TYPE, (SELECT COMMENTS FROM USER_COL_COMMENTS 
                     WHERE TABLE_NAME = '{table_name}' AND COLUMN_NAME = UTC.COLUMN_NAME) AS COMMENTS
                    FROM USER_TAB_COLUMNS UTC
                    WHERE TABLE_NAME='{table_name}'
                    ORDER BY COLUMN_ID
                """)
                columns = cursor.fetchall()
                cursor.close()
                
                info_part += "列信息:\n"
                for column in columns:
                    column_name, data_type, comments = column[0], column[1], column[2]
                    info_part += f"  - {column_name} ({data_type})"
                    if comments:
                        info_part += f": {comments}"
                    info_part += "\n"
                
                info_parts.append(info_part)
            
            return "\n".join(info_parts)
        except Exception as e:
            return f"获取表和列信息失败: {e}"

def get_sql_database():
    """
    获取适配的SQLDatabase对象，用于LangChain的SQL查询链
    """
    try:
        db = DMDatabase()
        return db
    except Exception as e:
        print(f"SQLDatabase适配器创建失败: {e}")
        return None

def clean_sql(sql_query):
    """清理SQL查询语句，移除格式问题"""
    # 移除可能的标记符号
    sql_query = sql_query.replace('```sql', '').replace('```', '')
    # 移除多余的空白字符
    sql_query = re.sub(r'\s+', ' ', sql_query).strip()
    # 移除反引号
    sql_query = sql_query.replace('`', '')
    # 移除可能的注释
    sql_query = re.sub(r'--.*', '', sql_query)
    # 特殊处理：移除可能导致解析错误的别名表达式（加强版）
    sql_query = re.sub(r'\[([A-Z_]+\.[A-Z_]+)\]', r'\1', sql_query)
    # 处理可能的简单别名表达式
    sql_query = re.sub(r'\[([A-Z_]+)\]', r'\1', sql_query)
    return sql_query

def create_sql_chain(chat_model=None):
    """
    创建SQL查询链
    """
    from langchain.chains import create_sql_query_chain
    
    # 获取SQLDatabase对象
    db = get_sql_database()
    
    if db is None:
        return None
    
    # 如果没有提供模型，使用默认模型
    if chat_model is None:
        chat_model = ChatOpenAI(
            model=os.getenv("OPENAI_MODEL", "deepseek-chat"),
            temperature=0
        )
    
    # 创建SQL查询链
    try:
        chain = create_sql_query_chain(chat_model, db)
        return chain
    except Exception as e:
        print(f"SQL查询链创建失败: {e}")
        return None

def execute_sql_query(query, db=None):
    """
    执行SQL查询
    """
    if db is None:
        db = get_sql_database()
        
    if db is None:
        return {"error": "数据库连接失败"}
    
    try:
        # 清理SQL查询语句
        query = clean_sql(query)
        # 执行查询
        result = db.run(query)
        return {"result": result}
    except Exception as e:
        return {"error": f"查询执行失败: {str(e)}"}

def analyze_user_question_and_get_relevant_tables(question, db=None):
    """
    分析用户问题并获取相关的表和列信息
    
    Args:
        question (str): 用户的自然语言问题
        db (DMDatabase): 数据库连接对象
    
    Returns:
        dict: 包含相关信息的字典
    """
    if db is None:
        db = get_sql_database()
    
    if db is None:
        return {"error": "数据库连接失败"}
    
    try:
        # 获取所有表和列信息
        all_table_info = db.get_all_table_and_column_info()
        
        return {
            "all_table_info": all_table_info,
            "status": "success"
        }
    except Exception as e:
        return {"error": f"分析用户问题失败: {str(e)}"}